package nlp;

import java.util.HashMap;
import java.util.HashSet;

import weka.classifiers.Classifier;
import weka.classifiers.Evaluation;
import weka.classifiers.bayes.NaiveBayes;
import weka.classifiers.functions.SMO;
import weka.classifiers.lazy.IBk;
import weka.classifiers.meta.AdaBoostM1;
import weka.classifiers.meta.LogitBoost;
import weka.classifiers.rules.PART;
import weka.classifiers.trees.DecisionStump;
import weka.classifiers.trees.J48;
import weka.classifiers.trees.J48graft;
import weka.core.Instances;

public class AuthorClassifier {
	public static String[] authorNames = { "Andrew.Lang.author.dir",
			"Charles.Kingsley.author.dir", "Charlotte.Mary.Yonge.author.dir",
			"G.K.Chesterton.author.dir", "H.G.Wells.author.dir",
			"Jacob.Abbott.author.dir", "John.Morley.author.dir",
			"John.Ruskin.author.dir", "R.M.Ballantyne.author.dir",
			"Robert.Louis.Stevenson.author.dir", "Samuel.Vaknin.author.dir",
			"Thomas.Carlyle.author.dir", "Thomas.Henry.Huxley.author.dir",
			"William.Dean.Howells.author.dir",
      "William.Henry.Giles.Kingston.author.dir", "Andrew.McCallum",
      "Christopher.Manning", "Lillian.Lee", "Raymond.Mooney",
      "Regina.Barzilay"};

	@SuppressWarnings("unchecked")
	public static void main(String[] argv) throws Exception {
		if (argv.length < 3) {
			System.err
					.println("Usage: <program> <train-data> <test-data> <classifier-class>");
			return;
		}

		// Create instances of required classifiers.
		HashMap<String, Classifier> classifierMap = new HashMap<String, Classifier>();
		classifierMap.put("smo", new SMO());
		classifierMap.put("bayes", new NaiveBayes());
		classifierMap.put("logitboost", new LogitBoost());
		classifierMap.put("adaboost", new AdaBoostM1());
		classifierMap.put("decisionstump", new DecisionStump());
		classifierMap.put("dtrees", new J48());
		classifierMap.put("dtreesgraft", new J48graft());
		classifierMap.put("knn", new IBk());
		classifierMap.put("rulebased", new PART());

		// Check classifier string.
		if (!classifierMap.containsKey(argv[2])) {
			System.err
					.println("Classifier not present: Use one of the following:");
			System.err.println("smo,bayes,logitboost,adaboost,decisionstump"
					+ "dtrees,dtreesgraft,knn,rulebased");
		}

		// Create the feature skips.
		HashSet<String> featureSkipSet = new HashSet<String>();
		if (argv.length == 4) {
			String[] features = argv[3].split(",");
			for (String feature : features) {
				featureSkipSet.add(feature);
			}
		}

		// Create training and test data.
		WekaFeatureManager featMgr = new WekaFeatureManager(argv[0], argv[1],
				featureSkipSet);
		Instances trainingSet = featMgr.GetTrainingInstances();
		Instances testingSet = featMgr.GetTestInstances();

		// Choose the classifier.
		Classifier classifier = classifierMap.get(argv[2]);
		classifier.buildClassifier(trainingSet);

		// Evaluate the classifier.
		Evaluation eval = new Evaluation(trainingSet);
		eval.evaluateModel(classifier, testingSet);

		// Report Summary.
		System.out.println(eval.toSummaryString());
		for (int i = 0; i < authorNames.length; ++i) {
			System.out.println(authorNames[i] + "," + eval.precision(i) + ","
					+ eval.recall(i) + "," + eval.fMeasure(i));
		}
	}
}
